use std::{mem::MaybeUninit, ops::Deref, sync::Arc, time::Duration};

use axum::{async_trait, extract::FromRequestParts, http::request::Parts};
use dashmap::DashMap;
use tokio::{
    sync::{mpsc, Mutex},
    time::timeout,
};

use crate::{
    error::{
        ec::{self},
        new_err_msg, AppErr, ErrorWrap,
    },
    server::{client::{
        message::{
            Message, MSG_TYPE_ACK, MSG_TYPE_NOTIFY, MSG_TYPE_NOTIFY_ACK, MSG_TYPE_PING,
            MSG_TYPE_PONG, MSG_TYPE_REQ, MSG_TYPE_RES,
        },
        msg_router::MessageRouter,
    }, header},
};



pub const ADV_NOTIFY: &'static str = "/adv/notify";

pub struct MessageHandler {
    msg_chan: Arc<MessageChannel>,
}

pub struct MessageChannel {
    tx: mpsc::Sender<Message>,
    rx: Mutex<mpsc::Receiver<Message>>,
    router: MessageRouter,
}

impl MessageChannel {
    fn new() -> Self {
        let (tx, rx) = mpsc::channel(10);
        Self {
            tx,
            rx: Mutex::new(rx),
            router: MessageRouter::new(),
        }
    }

    pub async fn poll(&self) -> Result<Message, AppErr> {
        let mut rx = timeout(Duration::from_secs(2), self.rx.lock())
            .await
            .wrap()?;
        let msg = timeout(Duration::from_secs(10), rx.recv())
            .await
            .map_err(|_| new_err_msg(ec::POLL_MSG_TIMEOUT, ec::POLL_MSG_TIMEOUT_MSG))?
            .wrap()?;
        Ok(msg)
    }

    pub fn post(&self, msg: Message) -> Result<(), AppErr> {
        self.router.post(msg)?;
        Ok(())
    }

    pub fn notify(&self, mut msg: Message) -> Result<(), AppErr> {
        msg.msg_type = MSG_TYPE_NOTIFY;
        self.tx.try_send(msg).wrap()?;
        Ok(())
    }

    pub async fn notify_ack(&self, mut msg: Message) -> Result<(), AppErr> {
        msg.seq = self.router.get_seq();
        msg.msg_type = MSG_TYPE_NOTIFY_ACK;
        let rx = self.router.expect_rx(msg.seq, MSG_TYPE_ACK);
        self.tx.send(msg).await.wrap()?;

        timeout(Duration::from_secs(3), rx)
            .await
            .map_err(|_| new_err_msg(ec::MSG_NO_ACK, "no recv ack"))?
            .wrap()?;
        Ok(())
    }

    pub async fn ping(&self) -> Result<(), AppErr> {
        let seq = self.router.get_seq();
        const PING: &'static str = "ping";
        let rx = self.router.expect_rx(seq, MSG_TYPE_PONG);
        let msg = Message::empty(seq, MSG_TYPE_PING, PING);
        self.tx.send(msg).await.wrap()?;

        timeout(Duration::from_secs(3), rx)
            .await
            .map_err(|_| new_err_msg(ec::MSG_NO_PONG, "no recv pong"))?
            .wrap()?;
        Ok(())
    }

    pub async fn request(&self, mut msg: Message, duration: Duration) -> Result<Message, AppErr> {
        msg.seq = self.router.get_seq();
        msg.msg_type = MSG_TYPE_REQ;
        let ack_rx = self.router.expect_rx(msg.seq, MSG_TYPE_ACK);
        let resp_rx = self.router.expect_rx(msg.seq, MSG_TYPE_RES);

        self.tx.send(msg).await.wrap()?;

        timeout(Duration::from_secs(3), ack_rx)
            .await
            .map_err(|_| new_err_msg(ec::MSG_NO_ACK, "no recv ack with req"))?
            .wrap()?;

        let resp = timeout(duration, resp_rx)
            .await
            .map_err(|_| new_err_msg(ec::MSG_NO_RESP, "no recv resp"))?
            .wrap()?;

        Ok(resp)
    }
}

impl Deref for MessageHandler {
    type Target = MessageChannel;

    fn deref(&self) -> &Self::Target {
        &self.msg_chan
    }
}

type Map = DashMap<u64, Arc<MessageChannel>>;

static mut DEVICE_INFO_MAP: MaybeUninit<Map> = MaybeUninit::uninit();

pub fn init() {
    unsafe {
        DEVICE_INFO_MAP.write(Map::new());
    }
}

pub fn save_device(id: u64) {
    get_map().insert(id, Arc::new(MessageChannel::new()));
}

#[inline(always)]
fn get_map() -> &'static Map {
    unsafe { DEVICE_INFO_MAP.assume_init_ref() }
}

fn get_message_handler(id: u64) -> Result<MessageHandler, AppErr> {
    let info = get_map()
        .get(&id)
        .ok_or(new_err_msg(
            ec::DEVICE_IS_NOT_LOGIN,
            ec::DEVICE_IS_NOT_LOGIN_MSG,
        ))?
        .clone();
    Ok(MessageHandler { msg_chan: info })
}

pub fn adv_notify_all() {
    let map = get_map();
    for item in map.iter() {
        _ = item
            .value()
            .notify(Message::empty(0, MSG_TYPE_NOTIFY, ADV_NOTIFY));
    }
}

impl MessageHandler {
    
    pub async fn notify_adv(&self) -> Result<(), AppErr> {
        self.notify_ack(Message::empty(0, MSG_TYPE_NOTIFY_ACK, ADV_NOTIFY)).await?;
        Ok(())
    }
}

#[async_trait]
impl<S: Send + Sync> FromRequestParts<S> for MessageHandler {
    type Rejection = AppErr;

    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
        let id: u64 = parts
            .headers
            .get(header::DEVICE_TOKEN)
            .ok_or(new_err_msg(ec::DEVICE_NO_TOKEN, ec::DEVICE_NO_TOKEN_MSG))?
            .to_str()
            .wrap()?
            .parse()
            .wrap()?;
        let info = get_message_handler(id)?;
        Ok(info)
    }
}
